"""
视频处理工具模块

该模块提供视频文件的读取、写入和处理功能。
主要功能：
1. 视频文件读取和解码
2. 视频帧处理和缓存
3. 视频写入和编码
4. 视频信息获取

作者：目标跟踪系统
日期：2024
"""

import cv2
import numpy as np
from typing import Optional, Tuple, Generator, List
import os
import threading
from queue import Queue


class VideoProcessor:
    """
    视频处理器类
    
    提供视频文件的读取、处理和写入功能。
    支持多种视频格式，具有帧缓存和多线程处理能力。
    
    参数：
        video_path: 视频文件路径
        buffer_size: 帧缓存大小，默认为30帧
    """
    
    def __init__(self, video_path: Optional[str] = None, buffer_size: int = 30):
        self.video_path = video_path
        self.buffer_size = buffer_size
        
        # 视频相关属性
        self.cap = None
        self.writer = None
        self.frame_count = 0
        self.fps = 0
        self.width = 0
        self.height = 0
        self.current_frame_index = 0
        
        # 缓存相关
        self.frame_buffer = Queue(maxsize=buffer_size)
        self.is_reading = False
        self.read_thread = None
        
        # 如果提供了视频路径，立即打开
        if video_path:
            self.open_video(video_path)
            
    def open_video(self, video_path: str) -> bool:
        """
        打开视频文件
        
        参数：
            video_path: 视频文件路径
            
        返回：
            是否成功打开视频
        """
        if not os.path.exists(video_path):
            print(f"错误：视频文件不存在 - {video_path}")
            return False
            
        # 关闭之前的视频
        self.close_video()
        
        # 打开新视频
        self.cap = cv2.VideoCapture(video_path)
        
        if not self.cap.isOpened():
            print(f"错误：无法打开视频文件 - {video_path}")
            return False
            
        # 获取视频信息
        self.video_path = video_path
        self.frame_count = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self.fps = self.cap.get(cv2.CAP_PROP_FPS)
        self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        self.current_frame_index = 0
        
        print(f"视频信息: {self.width}x{self.height}, {self.fps}fps, {self.frame_count}帧")
        return True
        
    def close_video(self) -> None:
        """
        关闭视频文件
        """
        # 停止读取线程
        self.stop_reading()
        
        # 关闭视频捕获
        if self.cap:
            self.cap.release()
            self.cap = None
            
        # 关闭视频写入
        if self.writer:
            self.writer.release()
            self.writer = None
            
    def read_frame(self) -> Tuple[bool, Optional[np.ndarray]]:
        """
        读取下一帧
        
        返回：
            (是否成功读取, 帧图像)
        """
        if not self.cap:
            return False, None
            
        ret, frame = self.cap.read()
        
        if ret:
            self.current_frame_index += 1
            
        return ret, frame
        
    def seek_frame(self, frame_index: int) -> bool:
        """
        跳转到指定帧
        
        参数：
            frame_index: 目标帧索引
            
        返回：
            是否成功跳转
        """
        if not self.cap:
            return False
            
        if frame_index < 0 or frame_index >= self.frame_count:
            return False
            
        self.cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
        self.current_frame_index = frame_index
        return True
        
    def get_frame_at(self, frame_index: int) -> Optional[np.ndarray]:
        """
        获取指定索引的帧
        
        参数：
            frame_index: 帧索引
            
        返回：
            帧图像，如果失败则返回None
        """
        if not self.seek_frame(frame_index):
            return None
            
        ret, frame = self.read_frame()
        return frame if ret else None
        
    def start_reading(self) -> None:
        """
        开始异步读取帧到缓存
        """
        if self.is_reading:
            return
            
        self.is_reading = True
        self.read_thread = threading.Thread(target=self._read_frames_to_buffer)
        self.read_thread.daemon = True
        self.read_thread.start()
        
    def stop_reading(self) -> None:
        """
        停止异步读取
        """
        self.is_reading = False
        
        if self.read_thread:
            self.read_thread.join(timeout=1.0)
            self.read_thread = None
            
        # 清空缓存
        while not self.frame_buffer.empty():
            try:
                self.frame_buffer.get_nowait()
            except:
                break
                
    def _read_frames_to_buffer(self) -> None:
        """
        异步读取帧到缓存的线程函数
        """
        while self.is_reading and self.cap:
            try:
                ret, frame = self.cap.read()
                if not ret:
                    break
                    
                # 如果缓存满了，丢弃最旧的帧
                if self.frame_buffer.full():
                    try:
                        self.frame_buffer.get_nowait()
                    except:
                        pass
                        
                self.frame_buffer.put((ret, frame), timeout=0.1)
                self.current_frame_index += 1
                
            except Exception as e:
                print(f"读取帧时发生错误: {e}")
                break
                
    def get_buffered_frame(self) -> Tuple[bool, Optional[np.ndarray]]:
        """
        从缓存获取帧
        
        返回：
            (是否成功获取, 帧图像)
        """
        try:
            return self.frame_buffer.get(timeout=0.1)
        except:
            return False, None
            
    def create_writer(self, output_path: str, fourcc: str = 'XVID') -> bool:
        """
        创建视频写入器
        
        参数：
            output_path: 输出视频路径
            fourcc: 视频编码器，默认为'XVID'
            
        返回：
            是否成功创建写入器
        """
        if not self.cap:
            print("错误：没有打开的视频文件")
            return False
            
        # 关闭之前的写入器
        if self.writer:
            self.writer.release()
            
        # 创建新的写入器
        fourcc_code = cv2.VideoWriter_fourcc(*fourcc)
        self.writer = cv2.VideoWriter(output_path, fourcc_code, self.fps, (self.width, self.height))
        
        if not self.writer.isOpened():
            print(f"错误：无法创建视频写入器 - {output_path}")
            return False
            
        return True
        
    def write_frame(self, frame: np.ndarray) -> bool:
        """
        写入一帧到输出视频
        
        参数：
            frame: 要写入的帧
            
        返回：
            是否成功写入
        """
        if not self.writer:
            return False
            
        try:
            self.writer.write(frame)
            return True
        except Exception as e:
            print(f"写入帧时发生错误: {e}")
            return False
            
    def get_video_info(self) -> dict:
        """
        获取视频信息
        
        返回：
            包含视频信息的字典
        """
        return {
            'path': self.video_path,
            'width': self.width,
            'height': self.height,
            'fps': self.fps,
            'frame_count': self.frame_count,
            'current_frame': self.current_frame_index,
            'duration': self.frame_count / self.fps if self.fps > 0 else 0
        }
        
    def resize_frame(self, frame: np.ndarray, target_size: Tuple[int, int]) -> np.ndarray:
        """
        调整帧大小
        
        参数：
            frame: 输入帧
            target_size: 目标尺寸 (width, height)
            
        返回：
            调整后的帧
        """
        return cv2.resize(frame, target_size)
        
    def extract_frames(self, start_frame: int = 0, end_frame: Optional[int] = None, 
                      step: int = 1) -> Generator[np.ndarray, None, None]:
        """
        提取指定范围的帧
        
        参数：
            start_frame: 起始帧索引
            end_frame: 结束帧索引，如果为None则到视频结尾
            step: 帧间隔
            
        返回：
            帧生成器
        """
        if not self.cap:
            return
            
        if end_frame is None:
            end_frame = self.frame_count
            
        # 跳转到起始帧
        self.seek_frame(start_frame)
        
        current_index = start_frame
        while current_index < end_frame:
            ret, frame = self.read_frame()
            if not ret:
                break
                
            if (current_index - start_frame) % step == 0:
                yield frame
                
            current_index += 1
            
    def __del__(self):
        """
        析构函数，确保资源被正确释放
        """
        self.close_video() 